#  -*- coding: utf-8 -*-

from pymongo import UpdateOne,ASCENDING, DESCENDING
from monitor.base_monitor import BaseMonitor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,multi_computer,get_code_name,get_trading_dates,get_diff_dates,get_sub_industry
from util.database import DB_CONN
import time
import pandas as pd
from datetime import datetime, timedelta
from factor.factor_module import FactorModule

"""
实现动量模型监控：每一日的最强板块跟踪
"""


class MomentumMonitor(BaseMonitor):
    def __init__(self):
        BaseMonitor.__init__(self, name='momentum')
        self.collection = DB_CONN['momentum']
        self.collection.create_index([('name', 1), ('date', 1),('type',1),('origin',1)])

    def monitoring(self, begin_date, end_date):
        dm = DataModule()
        fm = FactorModule()
        """
        计算指定时间范围内板块涨幅排名
        """

        dates = get_trading_dates(begin_date,end_date)

        for date in dates:
            start_time = time.time()
            update_requests = list()
            block_count_dict = dict()

            # Step1 取20日RS大于87的个股排名,rs_120>0代表上市超过半年（超过半年才有rs_120的数据）
            fm_date_df = fm.get_single_date_factors("rs", False, date)
            fm_date_df['sub_industry']= fm_date_df.apply(lambda row: get_sub_industry(row['code']), axis=1)
            #去掉ETF，没有对应的行业
            fm_date_df = fm_date_df[fm_date_df['sub_industry'] != ""]
            rs_df = fm_date_df.loc[(fm_date_df['rs_20'] >= 87) & (fm_date_df['rs_120'] >= 0)]
            #print(rs_df)

            #Step3 todo:去掉基金持股小于2%的个股

            # Step4 余下个股按板块排名
            groupby_fm_date_tuple = fm_date_df.groupby('sub_industry')
            for group in groupby_fm_date_tuple:
                block_count_dict.update(
                    {group[0]:group[1].index.size}
                )
            #print(block_count_dict)
            groupby_rs_tuple = rs_df.groupby('sub_industry')
            for group in groupby_rs_tuple:
                # Step5 计算动量分值，每天的分值记到Momentum表里，并记录是哪些个股
                code_list = list()
                name = group[0]
                group_df = group[1]
                block_cur_count = block_count_dict[name]
                momentum_value = round(group_df.index.size * group_df.index.size /block_cur_count,2)
                #print(name,group_df.index.size,block_cur_count,momentum_value)

                for index, row in group_df.iterrows():
                    code_list.append(row['code'])

                update_requests.append(
                    UpdateOne(
                        {'name': name, 'date': date, 'type': 'sub_industry','origin':"choice"},
                        {'$set': {
                            'name': name,
                            'date': date,
                            'type': 'sub_industry',
                            'origin': 'choice',
                            'momentum_value': momentum_value,
                            'code_list':code_list}},
                        upsert=True))

            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)

                end_time = time.time()
                print('填充动量数据,日期：%s，插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (date, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

        #Step6 计算排名并存入数据库
        self.calc_mom_rank(begin_date,end_date)

        return

    def calc_mom_rank(self,begin_date,end_date):
        dates = get_trading_dates(begin_date,end_date)
        for date in dates:
            start_time = time.time()
            update_requests = list()
            momentum_cursor = self.collection.find(
                {'date': date, 'type': 'sub_industry', "origin": 'choice'},
                projection={'code_list': False, '_id': False})

            if momentum_cursor.count() > 0:

                mom_df = pd.DataFrame([x for x in momentum_cursor])
                mom_df['rank'] = mom_df['momentum_value'].rank(ascending=False)

                for index,row in mom_df.iterrows():
                    update_requests.append(
                        UpdateOne(
                            {'name': row['name'], 'date': date, 'type': 'sub_industry', 'origin': "choice"},
                            {'$set': {'rank': int(row['rank'])}},
                            upsert=True))

                if len(update_requests) > 0:
                    update_result = self.collection.bulk_write(update_requests, ordered=False)

                    end_time = time.time()
                    print('填充动量排名,日期：%s，插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                          (date, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                          flush=True)
        return

    #Step7 对动量表进行分析，找出排名靠前，分值提升高，名次提升快等板块
    def analyze_momentum(self,begin_date,end_date):

        result_df = pd.DataFrame(columns=('date', 'Top1','Top1_v',
                                          'Top2', 'Top2_v',
                                          'Top3', 'Top3_v',
                                          'Top4', 'Top4_v',
                                          'Top5', 'Top5_v',
                                          'Top6', 'Top6_v',
                                          'Top7', 'Top7_v',
                                          'Top8', 'Top8_v',
                                          'Top9', 'Top9_v',
                                          'Top10', 'Top10_v',
                                          'Top11', 'Top11_v',
                                          'Top12', 'Top12_v',
                                          'Top13', 'Top13_v',
                                          'Top14', 'Top14_v',
                                          'Top15', 'Top15_v'))

        dates = get_trading_dates(begin_date,end_date)

        for date in dates:
            #print(date)
            date_result_dict = dict()
            date_result_dict.update({"date":date})
            momentum_cursor = self.collection.find(
                {'date': date, 'type': 'sub_industry', "origin": 'choice'},
                projection={'code_list': False, '_id': False})

            mom_df = pd.DataFrame([x for x in momentum_cursor])
            sort_df = mom_df.sort_values(by="momentum_value", ascending=False)
            top15_sort_df = sort_df.iloc[0:15]
            i = 1
            for index, row in top15_sort_df.iterrows():
                date_result_dict.update({f"Top{i}":row['name'],
                                         f"Top{i}_v": row['momentum_value']})
                i += 1
            result_df = result_df.append(date_result_dict, ignore_index=True)

        result_df.to_csv("momentum.csv")

        return

if __name__ == '__main__':
    # 执行因子的提取任务
    #hfq =HfqMAFactor()
    pd.set_option('display.width',500)
    pd.set_option('display.max_columns', 500)
    pd.set_option('display.max_colwidth', 500)
    #MomentumMonitor().monitoring('2021-02-19', '2021-02-19')
    current_date = datetime.now().strftime('%Y-%m-%d')
    MomentumMonitor().analyze_momentum('2021-01-01', current_date)
    #MomentumMonitor().calc_mom_rank('2017-01-01', '2021-02-19')
    momentum_cursor = DB_CONN.momentum.find(
        {'name':"化学原料",'type': 'sub_industry', "origin": 'choice'},
        projection={'code_list': False, '_id': False})
    mom_df = pd.DataFrame([x for x in momentum_cursor])
    mom_df.to_csv("化学原料.csv")